# file:Dataloader.py
"""导入必要的库"""
# 实现基本运算:
import torch
# 数据装载器:
from torch.utils.data import DataLoader as dataloader
# 读取数据集:
import torchvision.datasets as datasets
# 数据转换器:
import torchvision.transforms as transforms

transformer = transforms.Compose([
	# 转化成张量:
	transforms.ToTensor(),
	# 标准化数据:
	transforms.Normalize((0.1307,),(0.3081,))])

def get_train_data():
	"""获取训练数据"""
	# 下载训练数据:
	train_data = datasets.MNIST(root='./',
								download=True,
								train=True,
								transform=transformer)
	# 装载训练数据:
	train_loader = dataloader(train_data,
							  batch_size=64,
							  shuffle=True)
	return train_data, train_loader	  


def test_val_data():
	"""获取测试、验证数据"""
	# 下载数据
	t_v_data = datasets.MNIST(root='./',
							   download=True,
							   train=False,
							   transform=transformer)
	# 设置采样器(一半用于验证、一半用于测试):
	indices = range(len(t_v_data))
	test_sampler = torch.utils.data.sampler.SubsetRandomSampler(indices[:5000])
	val_sampler = torch.utils.data.sampler.SubsetRandomSampler(indices[5000:])
	# 装载数据:
	test_dataloader = dataloader(t_v_data,
								 sampler=test_sampler,
								 batch_size=64)
	val_dataloader = dataloader(t_v_data,
								sampler=val_sampler,
								batch_size=64)
	return t_v_data, test_dataloader, val_dataloader


